from typing import Dict

from ..utils import create_object
from .base_memory import Memory


class MemoryManager:
    """
    记忆管理器类
    
    负责管理不同会话的记忆实例，提供记忆的创建、获取、添加和重置功能
    """

    def __init__(self, cfg) -> None:
        """
        初始化记忆管理器
        
        参数:
            cfg: 记忆配置，用于创建记忆实例
        """
        self.cfg = cfg  # 记忆配置
        self.memory_map: Dict[str, Memory] = {}  # 会话ID到记忆实例的映射

    def create_instance(self, session_id):
        """
        创建新的记忆实例
        
        为指定会话ID创建一个新的记忆实例
        
        参数:
            session_id: 会话标识符
        """
        self.memory_map[session_id] = create_object(self.cfg)

    def get_memory(self, session_id=0, **kwargs) -> list:
        """
        获取记忆内容
        
        获取指定会话的记忆内容
        
        参数:
            session_id: 会话标识符，默认为0
            kwargs: 额外的参数
            
        返回:
            list: 记忆内容列表
        """
        return self.memory_map[session_id].get_memory(**kwargs)

    def add(self, memory, session_id=0, **kwargs) -> None:
        """
        添加记忆
        
        向指定会话添加新的记忆内容。如果会话不存在，会先创建新的会话。
        
        参数:
            memory: 要添加的记忆内容
            session_id: 会话标识符，默认为0
            kwargs: 额外的参数
        """
        # 如果会话不存在，创建新的会话实例
        if session_id not in self.memory_map:
            self.create_instance(session_id)
        self.memory_map[session_id].add(memory, **kwargs)

    def get(self, session_id=0) -> Memory:
        """
        获取记忆实例
        
        获取指定会话的记忆实例
        
        参数:
            session_id: 会话标识符，默认为0
            
        返回:
            Memory: 记忆实例，如果会话不存在则返回None
        """
        return self.memory_map.get(session_id, None)

    def reset(self, session_id=0) -> None:
        """
        重置记忆
        
        删除指定会话的记忆实例
        
        参数:
            session_id: 会话标识符，默认为0
        """
        if session_id in self.memory_map:
            del self.memory_map[session_id]
